import sys
import os
import cv2
import torch
from pathlib import Path
from PyQt5.QtWidgets import (QApplication, QMainWindow, QLabel, QPushButton,
                             QFileDialog, QVBoxLayout, QHBoxLayout, QWidget,
                             QProgressBar, QMessageBox, QComboBox, QGroupBox,
                             QSplitter, QSizePolicy)
from PyQt5.QtGui import QPixmap, QImage, QPalette, QColor
from PyQt5.QtCore import Qt, QThread, pyqtSignal

# 添加 YOLOv5 相关代码
FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))

# 导入 YOLOv5 相关模块
from models.common import DetectMultiBackend
from utils.general import (check_img_size, non_max_suppression,
                           scale_boxes, xyxy2xywh, increment_path)
from utils.plots import Annotator, colors
from utils.torch_utils import select_device


class YOLODetector:
    """YOLOv5 检测器封装类"""

    def __init__(self):
        self.device = select_device('')
        self.model = None
        self.names = None
        self.stride = None
        self.loaded = False

    def load_model(self, weights='yolov5s.pt'):
        """加载 YOLOv5 模型"""
        try:
            # 加载模型
            self.model = DetectMultiBackend(weights, device=self.device, dnn=False)
            self.stride, self.names, _ = self.model.stride, self.model.names, self.model.pt
            self.loaded = True
            return True, f"模型加载成功: {weights}"
        except Exception as e:
            return False, f"模型加载失败: {str(e)}"

    def detect_image(self, image_path, conf_thres=0.25, iou_thres=0.45, img_size=640, classes=None):
        """检测单张图片"""
        if not self.loaded:
            return None, "模型未加载"

        try:
            # 读取图片
            img0 = cv2.imread(image_path)
            if img0 is None:
                return None, "无法读取图片"

            # 预处理
            imgsz = check_img_size((img_size, img_size), s=self.stride)
            img = cv2.resize(img0, (imgsz, imgsz))
            img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB
            img = torch.from_numpy(img).to(self.device).float()
            img /= 255.0  # 归一化
            img = img[None]  # 扩展维度

            # 推理
            pred = self.model(img, augment=False)

            # NMS
            pred = non_max_suppression(pred, conf_thres, iou_thres, classes=classes, max_det=1000)

            # 处理结果
            det = pred[0]
            annotator = Annotator(img0.copy(), line_width=2, example=str(self.names))

            if len(det):
                # 调整框尺寸
                det[:, :4] = scale_boxes(img.shape[2:], det[:, :4], img0.shape).round()

                # 绘制检测结果
                for *xyxy, conf, cls in reversed(det):
                    c = int(cls)
                    label = f'{self.names[c]} {conf:.2f}'
                    annotator.box_label(xyxy, label, color=colors(c, True))

            result_img = annotator.result()
            return result_img, "检测成功"
        except Exception as e:
            return None, f"检测失败: {str(e)}"


class DetectionThread(QThread):
    """检测线程 - 用于后台处理"""
    progress = pyqtSignal(int, str)  # 进度，当前文件
    finished = pyqtSignal(str)  # 完成信号
    error = pyqtSignal(str)  # 错误信号

    def __init__(self, detector, source, output_dir, conf_thres, img_size, classes):
        super().__init__()
        self.detector = detector
        self.source = source
        self.output_dir = output_dir
        self.conf_thres = conf_thres
        self.img_size = img_size
        self.classes = classes
        self.canceled = False

    def run(self):
        try:
            # 创建输出目录
            save_dir = increment_path(Path(self.output_dir) / 'detection_results')
            save_dir.mkdir(parents=True, exist_ok=True)

            # 确定输入类型（文件或文件夹）
            if os.path.isfile(self.source):
                files = [self.source]
            else:
            # 获取所有支持的图片格式
                img_formats = ['bmp', 'jpg', 'jpeg', 'png', 'tif', 'tiff', 'dng', 'webp', 'mpo']
                files = [str(f) for f in Path(self.source).glob('*')
                         if f.suffix[1:].lower() in img_formats]

            total = len(files)
            processed = 0

            for i, img_path in enumerate(files):
                if self.canceled:
                    break

                # 更新进度
                self.progress.emit(int((i + 1) / total * 100), f"处理中: {Path(img_path).name}")

                # 检测图片
                result_img, status = self.detector.detect_image(
                    img_path, self.conf_thres, img_size=self.img_size, classes=self.classes
                )

                if result_img is None:
                    continue

                # 保存结果
                output_path = save_dir / Path(img_path).name
                cv2.imwrite(str(output_path), result_img)

            processed += 1

            if self.canceled:
                self.finished.emit(f"操作已取消，已处理 {processed}/{total} 张图片")
            else:
                self.finished.emit(f"完成! 已处理 {processed} 张图片，结果保存在: {save_dir}")

        except Exception as e:
            self.error.emit(f"处理出错: {str(e)}")


    def cancel(self):
        self.canceled = True


class ImageDetectionApp(QMainWindow):
    """主应用程序界面"""

    def __init__(self):
        super().__init__()
        self.detector = YOLODetector()
        self.detection_thread = None
        self.current_image_path = None
        self.init_ui()
        self.load_default_model()

    def init_ui(self):
        """初始化用户界面"""
        self.setWindowTitle('YOLOv5 图像检测系统')
        self.setGeometry(100, 100, 1200, 800)

        # 设置主布局
        main_widget = QWidget()
        main_layout = QVBoxLayout()
        main_widget.setLayout(main_layout)
        self.setCentralWidget(main_widget)

        # 顶部控制面板
        control_group = QGroupBox("控制面板")
        control_layout = QHBoxLayout()

        # 模型选择
        model_layout = QVBoxLayout()
        model_layout.addWidget(QLabel("选择模型:"))
        self.model_combo = QComboBox()
        self.model_combo.addItems(['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt'])
        model_layout.addWidget(self.model_combo)

        # 置信度阈值
        conf_layout = QVBoxLayout()
        conf_layout.addWidget(QLabel("置信度阈值:"))
        self.conf_combo = QComboBox()
        self.conf_combo.addItems(['0.15', '0.25', '0.35', '0.45', '0.55', '0.65', '0.75'])
        self.conf_combo.setCurrentText('0.25')
        conf_layout.addWidget(self.conf_combo)

        # 图像尺寸
        size_layout = QVBoxLayout()
        size_layout.addWidget(QLabel("图像尺寸:"))
        self.size_combo = QComboBox()
        self.size_combo.addItems(['320', '416', '512', '640', '768', '896', '1024'])
        self.size_combo.setCurrentText('640')
        size_layout.addWidget(self.size_combo)

        # 按钮区域
        button_layout = QVBoxLayout()
        self.load_model_btn = QPushButton("加载模型")
        self.load_model_btn.clicked.connect(self.load_model)
        button_layout.addWidget(self.load_model_btn)

        self.select_img_btn = QPushButton("选择图片")
        self.select_img_btn.clicked.connect(self.select_image)
        button_layout.addWidget(self.select_img_btn)

        self.select_dir_btn = QPushButton("选择文件夹")
        self.select_dir_btn.clicked.connect(self.select_directory)
        button_layout.addWidget(self.select_dir_btn)

        self.detect_btn = QPushButton("开始检测")
        self.detect_btn.clicked.connect(self.start_detection)
        self.detect_btn.setEnabled(False)
        button_layout.addWidget(self.detect_btn)

        # 添加到控制面板
        control_layout.addLayout(model_layout)
        control_layout.addLayout(conf_layout)
        control_layout.addLayout(size_layout)
        control_layout.addLayout(button_layout)
        control_layout.addStretch()
        control_group.setLayout(control_layout)

        # 图像显示区域
        image_group = QGroupBox("图像预览")
        image_layout = QHBoxLayout()

        # 使用QSplitter实现可调整大小的分割
        splitter = QSplitter(Qt.Horizontal)

        # 原始图像区域
        self.original_label = QLabel()
        self.original_label.setAlignment(Qt.AlignCenter)
        self.original_label.setMinimumSize(400, 400)
        self.original_label.setStyleSheet("background-color: #2D2D2D;")
        self.original_label.setText("原始图像")
        splitter.addWidget(self.original_label)

        # 检测结果区域
        self.result_label = QLabel()
        self.result_label.setAlignment(Qt.AlignCenter)
        self.result_label.setMinimumSize(400, 400)
        self.result_label.setStyleSheet("background-color: #2D2D2D;")
        self.result_label.setText("检测结果")
        splitter.addWidget(self.result_label)

        # 设置分割比例
        splitter.setSizes([600, 600])
        image_layout.addWidget(splitter)
        image_group.setLayout(image_layout)

        # 进度条
        progress_group = QGroupBox("进度")
        progress_layout = QVBoxLayout()

        self.progress_bar = QProgressBar()
        self.progress_bar.setRange(0, 100)
        self.progress_bar.setTextVisible(True)

        self.status_label = QLabel("就绪")
        self.status_label.setAlignment(Qt.AlignCenter)

        self.cancel_btn = QPushButton("取消检测")
        self.cancel_btn.clicked.connect(self.cancel_detection)
        self.cancel_btn.setEnabled(False)

        progress_layout.addWidget(self.progress_bar)
        progress_layout.addWidget(self.status_label)
        progress_layout.addWidget(self.cancel_btn)
        progress_group.setLayout(progress_layout)

        # 添加到主布局
        main_layout.addWidget(control_group)
        main_layout.addWidget(image_group, 1)  # 设置伸缩因子为1
        main_layout.addWidget(progress_group)

        # 设置样式
        self.setStyleSheet("""
            QGroupBox {
                font-weight: bold;
                border: 1px solid gray;
                border-radius: 5px;
                margin-top: 1ex;
                padding-top: 10px;
            }
            QGroupBox::title {
                subcontrol-origin: margin;
                left: 10px;
                padding: 0 5px;
            }
            QLabel {
                font-size: 12px;
            }
            QPushButton {
                padding: 5px;
                font-weight: bold;
            }
            QProgressBar {
                height: 20px;
                text-align: center;
            }
        """)

        # 设置深色主题
        dark_palette = QPalette()
        dark_palette.setColor(QPalette.Window, QColor(53, 53, 53))
        dark_palette.setColor(QPalette.WindowText, Qt.white)
        dark_palette.setColor(QPalette.Base, QColor(35, 35, 35))
        dark_palette.setColor(QPalette.AlternateBase, QColor(53, 53, 53))
        dark_palette.setColor(QPalette.ToolTipBase, Qt.white)
        dark_palette.setColor(QPalette.ToolTipText, Qt.white)
        dark_palette.setColor(QPalette.Text, Qt.white)
        dark_palette.setColor(QPalette.Button, QColor(53, 53, 53))
        dark_palette.setColor(QPalette.ButtonText, Qt.white)
        dark_palette.setColor(QPalette.BrightText, Qt.red)
        dark_palette.setColor(QPalette.Highlight, QColor(142, 45, 197).lighter())
        dark_palette.setColor(QPalette.HighlightedText, Qt.black)
        self.setPalette(dark_palette)

    def load_default_model(self):
        """加载默认模型"""
        success, message = self.detector.load_model()
        if success:
            self.status_label.setText(message)
        else:
            QMessageBox.critical(self, "错误", message)

    def load_model(self):
        """加载选定的模型"""
        model_name = self.model_combo.currentText()
        success, message = self.detector.load_model(model_name)

        if success:
            self.status_label.setText(message)
            self.detect_btn.setEnabled(True)
        else:
            QMessageBox.critical(self, "模型加载错误", message)

    def select_image(self):
        """选择单张图片"""
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择图片", "",
            "图片文件 (*.jpg *.jpeg *.png *.bmp *.tif *.tiff)"
        )

        if file_path:
            self.current_image_path = file_path
            self.show_image(self.original_label, file_path)
            self.result_label.clear()
            self.result_label.setText("检测结果")
            self.status_label.setText(f"已选择图片: {Path(file_path).name}")
            self.detect_btn.setEnabled(True)

    def select_directory(self):
        """选择文件夹"""
        dir_path = QFileDialog.getExistingDirectory(
            self, "选择文件夹", ""
        )

        if dir_path:
            self.current_image_path = dir_path
            self.original_label.clear()
            self.original_label.setText("已选择文件夹")
            self.result_label.clear()
            self.result_label.setText("点击开始检测处理整个文件夹")
            self.status_label.setText(f"已选择文件夹: {Path(dir_path).name}")
            self.detect_btn.setEnabled(True)

    def show_image(self, label, image_path):
        """在标签上显示图片"""
        pixmap = QPixmap(image_path)
        if not pixmap.isNull():
            # 缩放图片以适应标签，同时保持宽高比
            label.setPixmap(pixmap.scaled(
                label.width(), label.height(),
                Qt.KeepAspectRatio, Qt.SmoothTransformation
            ))

    def start_detection(self):
        """开始检测处理"""
        if not self.current_image_path:
            QMessageBox.warning(self, "警告", "请先选择图片或文件夹")
            return

        # 如果是单张图片
        if os.path.isfile(self.current_image_path):
            conf_thres = float(self.conf_combo.currentText())
            img_size = int(self.size_combo.currentText())

            # 检测图片
            result_img, status = self.detector.detect_image(
                self.current_image_path, conf_thres, img_size
            )

            if result_img is None:
                QMessageBox.critical(self, "错误", status)
                return

            # 显示结果
            height, width, channel = result_img.shape
            bytes_per_line = 3 * width
            q_img = QImage(result_img.data, width, height, bytes_per_line, QImage.Format_RGB888).rgbSwapped()
            pixmap = QPixmap.fromImage(q_img)
            self.result_label.setPixmap(pixmap.scaled(
                self.result_label.width(), self.result_label.height(),
                Qt.KeepAspectRatio, Qt.SmoothTransformation
            ))

            self.status_label.setText("单张图片检测完成")
            QMessageBox.information(self, "完成", "图片检测完成!")

        # 如果是文件夹
        else:
            # 创建输出目录
            output_dir = QFileDialog.getExistingDirectory(
                self, "选择结果保存位置", ""
            )

            if not output_dir:
                return

            # 获取参数
            conf_thres = float(self.conf_combo.currentText())
            img_size = int(self.size_combo.currentText())

            # 禁用按钮
            self.select_img_btn.setEnabled(False)
            self.select_dir_btn.setEnabled(False)
            self.detect_btn.setEnabled(False)
            self.cancel_btn.setEnabled(True)

            # 创建并启动检测线程
            self.detection_thread = DetectionThread(
                self.detector,
                self.current_image_path,
                output_dir,
                conf_thres,
                img_size,
                classes=None
            )

            # 连接信号
            self.detection_thread.progress.connect(self.update_progress)
            self.detection_thread.finished.connect(self.detection_finished)
            self.detection_thread.error.connect(self.detection_error)

            self.detection_thread.start()

    def update_progress(self, value, message):
        """更新进度条和状态"""
        self.progress_bar.setValue(value)
        self.status_label.setText(message)

    def detection_finished(self, message):
        """检测完成处理"""
        self.progress_bar.setValue(0)
        self.status_label.setText(message)

        # 启用按钮
        self.select_img_btn.setEnabled(True)
        self.select_dir_btn.setEnabled(True)
        self.detect_btn.setEnabled(True)
        self.cancel_btn.setEnabled(False)

        QMessageBox.information(self, "完成", message)

    def detection_error(self, message):
        """检测错误处理"""
        self.progress_bar.setValue(0)
        self.status_label.setText(message)

        # 启用按钮
        self.select_img_btn.setEnabled(True)
        self.select_dir_btn.setEnabled(True)
        self.detect_btn.setEnabled(True)
        self.cancel_btn.setEnabled(False)

        QMessageBox.critical(self, "错误", message)

    def cancel_detection(self):
        """取消检测"""
        if self.detection_thread and self.detection_thread.isRunning():
            self.detection_thread.cancel()
            self.status_label.setText("正在取消操作...")

    def resizeEvent(self, event):
        """窗口大小改变事件 - 重新缩放图片"""
        super().resizeEvent(event)

        # 重新显示原始图片
        if self.current_image_path and os.path.isfile(self.current_image_path):
            self.show_image(self.original_label, self.current_image_path)

        # 重新显示结果图片
        if hasattr(self.result_label, 'pixmap') and self.result_label.pixmap():
            pixmap = self.result_label.pixmap()
            scaled_pixmap = pixmap.scaled(
                self.result_label.width(), self.result_label.height(),
                Qt.KeepAspectRatio, Qt.SmoothTransformation
            )
            self.result_label.setPixmap(scaled_pixmap)


if __name__ == '__main__':
    app = QApplication(sys.argv)

    # 检查CUDA是否可用
    if torch.cuda.is_available():
        device_name = torch.cuda.get_device_name(0)
        app.setApplicationName(f"YOLOv5 图像检测系统 (GPU: {device_name})")
    else:
        app.setApplicationName("YOLOv5 图像检测系统 (CPU)")

    window = ImageDetectionApp()
    window.show()
    sys.exit(app.exec_())